Skip to content

[Feat] Add Sampled-token and top-k OPD Support to RL Training - #1977

Open
YanhuiDua wants to merge 3 commits into
InternLM:mainfrom
YanhuiDua:dev-opd
Open

[Feat] Add Sampled-token and top-k OPD Support to RL Training #1977
YanhuiDua wants to merge 3 commits into
InternLM:mainfrom
YanhuiDua:dev-opd

Conversation

@YanhuiDua

@YanhuiDua YanhuiDua commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

背景

本 PR 为 XTuner 引入统一的 Online Policy Distillation(OPD)架构,主要目标包括:

  • 同时支持 sampled-token OPD 和 teacher-selected Top-K OPD。
  • 同时支持单 Teacher 和多 Teacher。
  • sampled-token 与 Top-K Teacher target 均支持两种生产方式:
    • 由外部推理引擎计算;
    • 由训练引擎内的 frozen Teacher 计算。
    • 将 PG-OPD 的 advantage 构造和直接反向传播的蒸馏目标统一收敛到 distillation_loss.py,避免将 OPD 数学逻辑分散在 AgentLoop、Trainer 或 GRPO loss 中.

以下能力暂不包含在本 PR 范围内:

  • student-selected Top-K。
  • 对任意 token ID 进行外部 Teacher 精确打分。
  • Full-vocabulary distillation。
  • 在同一次训练中混用 RolloutTeacherConfig 和 TrainTeacherConfig。

LMDeploy推理依赖PR: InternLM/lmdeploy#4793

主要修改

配置设计

蒸馏 loss 通过 DistillationLossConfig 统一配置:

  loss_cfg = DistillationLossConfig(
      policy_loss_cfg=policy_loss_cfg,
      loss_mode="k1",  # k1, forward, reverse, forward_kl_topk
      use_policy_gradient=True,
      task_adv_weight=0.0,
      distillation_loss_weight=1.0,
      top_k=64,
      log_prob_min_clamp=-10.0,
      loss_max_clamp=10.0,
  )

其中:

  • use_policy_gradient 控制是否同时启用 policy-gradient loss,目前仅支持 sampled-token OPD 与 PG loss 组合。
  • log_prob_min_clamp 用于限制 selected-token log probability 的下界。
  • loss_max_clamp 用于截断逐 token 的 distillation loss。
  • top_k 用于配置 teacher-selected Top-K 的候选数量。

外部推理引擎 Teacher

distillation_config = DistillationConfig(
    loss_config=loss_cfg,
    teachers=[
        RolloutTeacherConfig(
            name="teacher",
            launch_config=RolloutTeacherLaunchConfig(
                model_path=teacher_model_path,
                server_port=13141,
            ),
        ),
    ],
    data_source_teacher_map={
        "math": "teacher",
    },
)

训练引擎内 Frozen Teacher

distillation_config = DistillationConfig(
    loss_config=loss_cfg,
    teachers=[
        TrainTeacherConfig(
            name="teacher",
            model_path=teacher_model_path,
            model_cfg=teacher_model_cfg,
            fsdp_cfg=teacher_fsdp_cfg,
        ),
    ],
    data_source_teacher_map={
        "math": "teacher",
    },
)

同一次训练仅允许选择一种 Teacher runtime:外部推理引擎或训练引擎内 frozen Teacher。

Rollout流程

image

Train流程

image

@YanhuiDua YanhuiDua changed the title [Feat] Add PG-OPD Support to RL Training [Feat] Add Sampled-token OPD Support to RL Training Jul 27, 2026
@YanhuiDua YanhuiDua changed the title [Feat] Add Sampled-token OPD Support to RL Training [Feat] Add Sampled-token OPD Support with SGLang to RL Training Jul 27, 2026
@YanhuiDua YanhuiDua changed the title [Feat] Add Sampled-token OPD Support with SGLang to RL Training [Feat] Add Sampled-token OPD Support to RL Training Jul 28, 2026
@YanhuiDua

Copy link
Copy Markdown
Collaborator Author

@claude review

@YanhuiDua YanhuiDua changed the title [Feat] Add Sampled-token OPD Support to RL Training [Feat] Add Sampled-token and top-k OPD Support to RL Training Aug 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant